# -*- coding: UTF-8 -*-
from PIL import Image
import os
import torch
from torchvision.transforms import Compose, ToTensor, ToPILImage, functional

from torchvision.utils import save_image
from util.LBAMModel import LBAMModel
from util.runtimeError import runtimeError

ImageTransform = Compose([
    ToTensor(),
])
pretrainPath = ""


def repairImage(filepaths, rootpath):
    try:
        pathlist = filepaths.split('\n\n')
        pathlist[-1] = pathlist[-1].rstrip('\n')
        for i in range(len(pathlist)):
            pathlist[i] = pathlist[i].split('->')[1]
        for index, path in enumerate(pathlist, 1):
<<<<<<< develop
            image_pil = Image.open(path).convert('RGB').resize((256, 256), resample=Image.BICUBIC)
            image = ImageTransform(image_pil)
            # 得到掩膜情况
            mask = getMaskImage(image_pil)
            mask = 1 - mask
            image = image * mask
            inputImage = torch.cat((image, mask[0].unsqueeze(0)), 0)
            inputImage = inputImage.unsqueeze(0)
            mask = mask.unsqueeze(0)
            netG = LBAMModel(4, 3)
            netG.load_state_dict(torch.load(rootpath + '/data/LBAM_NoBN_ParisStreetView.pth'))
            for param in netG.parameters():
                param.requires_grad = False
            netG.eval()
            # print(netG.reverseConv5.updateMask.alpha)
            if torch.cuda.is_available():
                netG = netG.cuda()
                inputImage = inputImage.cuda()
                mask = mask.cuda()
            else:
                netG = netG.cpu()
                inputImage = inputImage.cpu()
                mask = mask.cpu()

            output = netG(inputImage, mask)
            output = output * (1 - mask) + inputImage[:, 0:3, :, :] * mask
            if not os.path.exists(rootpath + '/output/inpainting/'):
                os.makedirs(rootpath + '/output/inpainting/')
            save_image(output, rootpath + '/output/inpainting/' + str(index) + '.png')
    except Exception as e:
        if e.args[0].find('CUDA out of memory') != -1:
            raise runtimeError("处理图片：" + str(index) + " 遇到问题。显存不够。具体：" + e.args[0])
        else:
            raise runtimeError("处理图片：" + str(index) + " 遇到问题。具体：" + e.args[0])


def getMaskImage(image):
    mask = functional.to_grayscale(image, 3)
    mask = ImageTransform(mask)
    zeros = mask < 1
    mask.masked_fill_(zeros, 0)
#    maskimg = ToPILImage()(mask.data.cpu())
#    maskimg.save("mask.png")
    return mask
